"""
RISC-V代码生成器工具函数

提供代码生成过程中需要的各种辅助功能。
"""

class LabelManager:
    """标签管理器，用于生成和管理跳转标签"""
    
    def __init__(self):
        self.counter = 0
        self.labels = {}
    
    def generate_label(self, prefix="L"):
        """生成唯一的标签名"""
        self.counter += 1
        return f"{prefix}{self.counter}"
    
    def create_loop_labels(self):
        """创建循环使用的标签对"""
        loop_start = self.generate_label("loop_start_")
        loop_end = self.generate_label("loop_end_")
        return loop_start, loop_end
    
    def create_if_labels(self):
        """创建条件判断使用的标签"""
        if_else = self.generate_label("if_else_")
        if_end = self.generate_label("if_end_")
        return if_else, if_end


class CodeSection:
    """代码段，用于组织不同部分的代码"""
    
    def __init__(self, name):
        self.name = name
        self.instructions = []
        self.labels = set()
    
    def add_instruction(self, instruction):
        """添加指令"""
        self.instructions.append(instruction)
    
    def add_label(self, label):
        """添加标签"""
        self.labels.add(label)
        self.instructions.append(f"{label}:")
    
    def clear(self):
        """清空代码段"""
        self.instructions.clear()
        self.labels.clear()
    
    def __str__(self):
        return '\n'.join(self.instructions)


class StackManager:
    """栈管理器，用于管理函数调用栈和局部变量"""
    
    def __init__(self, reg_manager):
        self.reg_manager = reg_manager
        self.stack_offset = 0
        self.saved_registers = []
    
    def allocate_stack_space(self, size):
        """分配栈空间"""
        self.stack_offset += size
        return -self.stack_offset
    
    def save_register(self, code_gen, reg):
        """保存寄存器到栈"""
        offset = self.allocate_stack_space(4)  # 32位系统，每个寄存器4字节
        code_gen.instructions.sw(reg, offset, 'sp')
        self.saved_registers.append((reg, offset))
        return offset
    
    def restore_register(self, code_gen, reg):
        """从栈恢复寄存器"""
        for saved_reg, offset in self.saved_registers:
            if saved_reg == reg:
                code_gen.instructions.lw(reg, offset, 'sp')
                self.saved_registers.remove((saved_reg, offset))
                return
        raise ValueError(f"Register {reg} not found in saved registers")
    
    def restore_all_registers(self, code_gen):
        """恢复所有保存的寄存器"""
        for reg, offset in self.saved_registers:
            code_gen.instructions.lw(reg, offset, 'sp')
        self.saved_registers.clear()
    
    def setup_function_prologue(self, code_gen):
        """设置函数序言"""
        # 保存返回地址和帧指针
        code_gen.instructions.addi('sp', 'sp', -8)
        code_gen.instructions.sw('ra', 4, 'sp')
        code_gen.instructions.sw('s0', 0, 'sp')
        code_gen.instructions.addi('s0', 'sp', 8)
    
    def setup_function_epilogue(self, code_gen):
        """设置函数结语"""
        # 恢复返回地址和帧指针
        code_gen.instructions.lw('ra', 4, 'sp')
        code_gen.instructions.lw('s0', 0, 'sp')
        code_gen.instructions.addi('sp', 'sp', 8)
        code_gen.instructions.ret()


def format_immediate(value):
    """格式化立即数，确保在RISC-V允许的范围内"""
    if isinstance(value, str):
        return value
    
    # RISC-V立即数范围检查
    if -2048 <= value <= 2047:
        return str(value)
    else:
        # 对于大的立即数，需要使用lui+addi组合
        # 这里返回原值，让调用者处理
        return str(value)


def is_power_of_2(n):
    """检查一个数是否是2的幂"""
    return n > 0 and (n & (n - 1)) == 0


def log2(n):
    """计算以2为底的对数"""
    if not is_power_of_2(n):
        raise ValueError(f"{n} is not a power of 2")
    
    result = 0
    while n > 1:
        n >>= 1
        result += 1
    return result


def optimize_constant_multiplication(value):
    """优化常数乘法，将乘法转换为移位和加法"""
    optimizations = []
    
    if value == 0:
        return [('li', 0)]
    elif value == 1:
        return [('mv',)]
    elif is_power_of_2(value):
        return [('slli', log2(value))]
    elif value == 3:
        # x * 3 = x + (x << 1)
        return [('slli_add', 1)]
    elif value == 5:
        # x * 5 = x + (x << 2)
        return [('slli_add', 2)]
    elif value == 9:
        # x * 9 = x + (x << 3)
        return [('slli_add', 3)]
    else:
        return [('mul',)] 